"""
划分训练集，测试集，验证集
"""
import os
import shutil
from tqdm import tqdm
import random

origin_path = "/media/liyan/3b75ab20-92d3-4be6-81c2-4e1798e2fe16/我的数据集/图像分类数据集/总/"
train_path = "/media/liyan/3b75ab20-92d3-4be6-81c2-4e1798e2fe16/我的数据集/图像分类数据集/new1/train/"
val_path = "/media/liyan/3b75ab20-92d3-4be6-81c2-4e1798e2fe16/我的数据集/图像分类数据集/new1/val/"
test_path = "/media/liyan/3b75ab20-92d3-4be6-81c2-4e1798e2fe16/我的数据集/图像分类数据集/new1/test/"

if True:
    class_directory = os.listdir(origin_path)
    for class_name in class_directory:
        os.makedirs(train_path+class_name)
        os.makedirs(val_path+class_name)
        os.makedirs(test_path+class_name)

    for direc in tqdm(class_directory,ncols=150):
        all_files = os.listdir(origin_path+direc)
        random.shuffle(all_files)
        length = len(all_files)

        train_data_path = [origin_path+direc+os.sep+filename for filename in all_files[:int(length*0.8)]]
        val_data_path = [origin_path+direc+os.sep+filename for filename in all_files[int(length*0.8):int(length*0.9)]]
        test_data_path = [origin_path+direc+os.sep+filename for filename in all_files[int(length*0.9):]]

        for path in train_data_path:
            dist = train_path+direc
            shutil.copy(path,dist)
        for path in val_data_path:
            dist = val_path+direc
            shutil.copy(path,dist)
        for path in test_data_path:
            dist = test_path+direc
            shutil.copy(path,dist)


